import socket
import struct
import json
from threading import Thread
import time

class MulticastReceiver:
    """
    UDP组播数据接收器
    """
    
    def __init__(self, multicast_group, port, callback=None, interface_ip=None):
        """
        初始化接收器
        :param multicast_group: 组播地址 (如 '239.255.255.250')
        :param port: 监听端口
        :param callback: 数据回调函数
        :param interface_ip: 指定网络接口的IP地址
        """
        self.multicast_group = multicast_group
        self.port = port
        self.callback = callback
        self.interface_ip = interface_ip
        self.running = False
        self.socket = None
        self.thread = None
    
    def start(self):
        """启动接收器"""
        try:
            # 创建UDP套接字
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            
            # 设置超时时间，避免recvfrom永久阻塞
            self.socket.settimeout(1.0)
            
            # 绑定到所有接口
            self.socket.bind(('', self.port))
            print(f"绑定到端口: {self.port}")
            
            # 加入多播组 - 使用指定接口
            if self.interface_ip:
                # 指定网络接口
                group = socket.inet_aton(self.multicast_group)
                iface = socket.inet_aton(self.interface_ip)
                mreq = group + iface
                self.socket.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)
                print(f"加入组播组 {self.multicast_group}，使用接口 {self.interface_ip}")
            else:
                # 使用默认接口
                mreq = struct.pack('4sl', socket.inet_aton(self.multicast_group), socket.INADDR_ANY)
                self.socket.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)
                print(f"加入组播组 {self.multicast_group}，使用默认接口")
            
            self.running = True
            # 启动监听线程
            self.thread = Thread(target=self.__listen, daemon=True)
            self.thread.start()
            print(f"启动接收器成功，监听组播地址: {self.multicast_group}:{self.port}")
            # return True
        except Exception as e:
            # print(f"启动接收器失败: {e}")
            raise e
            # return False
    
    def stop(self):
        """停止接收器"""
        self.running = False
        if self.socket:
            self.socket.close()
        if self.thread and self.thread.is_alive():
            self.thread.join(timeout=1.0)
    
    def __listen(self):
        """监听数据"""
        while self.running:
            try:
                data, addr = self.socket.recvfrom(4096)
                if self.callback:
                    # 原始数据回调
                    self.callback(data, addr)
            except socket.timeout:
                # 超时是正常的，继续循环
                continue
            except Exception as e:
                if self.running:  # 忽略正常停止导致的异常
                    print(f"接收数据异常: {e}")
    
    def __enter__(self):
        self.start()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stop()


def test_callback(data, addr):
    """测试回调函数"""
    print(f"来自 {addr} 的数据: {data.decode('utf-8', errors='ignore')}")

if __name__ == '__main__':
    # 测试用例
    MULTICAST_GROUP = '224.1.1.100'
    PORT = 5555
    
    # 使用上下文管理器确保资源正确释放
    with MulticastReceiver(MULTICAST_GROUP, PORT, test_callback) as receiver:
        print("正在监听组播数据...")
        try:
            # 保持主线程运行
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            print("\n用户中断，停止接收")
